package com.spring.demo.controller;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

@RestController
@RequestMapping("/sse")
public class SseController {

    private static final Logger log = LoggerFactory.getLogger(SseController.class);

    @RequestMapping("/emitter")
    public SseEmitter sse(SseEmitter emitter) {
        AtomicInteger atomicInteger = new AtomicInteger();
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(() -> {
            try {
                int i = atomicInteger.getAndIncrement();
                if (i < 5) {
                    emitter.send(i);
                } else {
                    emitter.complete();
                    executor.shutdown();
                }
            } catch (IOException e) {
                emitter.completeWithError(e);
            }
        }, 0, 1, TimeUnit.SECONDS);

        return emitter;
    }

    @RequestMapping("/response")
    public void response(HttpServletResponse response) {
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        try (PrintWriter writer = response.getWriter()) {
            for (int i = 0; i < 5; i++) {
                if (i > 0) {
                    Thread.sleep(1000);
                }
                writer.write("data:" + i + "\n\n");
                writer.flush();
            }
        } catch (Exception e) {
            log.info("error", e);
        }
    }
}
